from fastapi import WebSocket
import base64
from typing import Dict
from biz.utils.logger import logger
from biz.integrations.asr.voice_assistant import VoiceAssistant

active_connections: Dict[str, WebSocket] = {}
active_assistants: Dict[str, VoiceAssistant] = {}

async def handle_websocket(websocket: WebSocket):
    await websocket.accept()
    client_id = id(websocket)
    active_connections[client_id] = websocket
    voice_assistant = VoiceAssistant(websocket=websocket)
    active_assistants[client_id] = voice_assistant
    logger.info(f"新的WebSocket连接: {client_id}")
    # 连接建立后立即发送client_id给客户端
    await websocket.send_json({
        "status": "connected", 
        "client_id": client_id,
        "message": "WebSocket连接已建立"
    })
    try:
        while True:
            data = await websocket.receive_text()
            try:
                audio_chunk = base64.b64decode(data)
                result = await voice_assistant.process_audio(bytes(audio_chunk))
                await websocket.send_json(result)
            except Exception as e:
                logger.error(f"处理WebSocket消息时发生错误: {str(e)}")
                await websocket.send_json({"status": "error", "message": str(e)})
    
    except Exception as e:
        logger.error(f"WebSocket连接错误: {str(e)}")
    finally:
        # 先主动关闭连接
        try:
            await voice_assistant._asr_close()
        except Exception as e:
            logger.error(f"关闭ASR WebSocket连接时发生错误: {str(e)}")
        
        # 然后从字典中删除引用
        del active_connections[client_id]
        del active_assistants[client_id] 
        logger.info(f"WebSocket连接已从活跃连接列表中移除: {client_id}")

def get_active_assistants():
    return active_assistants